import pandas as pd
import numpy as np

def calculate_metrics(df, selected_years=None, dimension=None):
    """
    计算差旅费分析指标
    
    Parameters:
    df: DataFrame - 原始数据
    selected_years: list - 选择的年度列表
    dimension: list - 选择的维度 '同城市' 或 '同一酒店'
    month_year: list - 选择的级别列表 ['月度', '年度']
    """
    # 创建原始数据的副本
    result = df.copy()

    # 剔除住宿天数为0和发票金额为0的
    result = result[(result["住宿天数"] != 0) & (result["发票含税金额"] != 0)]

    result_1 = pd.read_excel("data/差旅报销行程明细(商旅).xlsx")
    result_1["城市名称"] = result_1["目的地名称"]
    result_1["发票人员名称"] = result_1["人员名称"]

    result = result.merge(result_1[["唯一标识", "城市名称", "出发时间"]], on = ["唯一标识", "城市名称"], how = "left")
    result = result.merge(result_1[["唯一标识", "发票人员名称", "报销级别名称"]], on=["唯一标识", "发票人员名称"], how="left")
    result["出发时间"] = result["出发时间"].apply(lambda x:str(x)[:7])
    result.rename(columns={"出发时间":"月度", "报销级别名称":"职级"}, inplace=True)

    # 过滤选定年度的数据
    if selected_years:
        result = result[result['年度'].isin(selected_years)]

    # 计算同城市均值
    city_means = result.groupby(['年度', '城市名称', '职级']).apply(
        lambda x: x['发票含税金额'].sum() / x['住宿天数'].sum()
    ).reset_index()
    city_means.columns = ['年度', '城市名称', '职级', '同年度同城市均值']

    # 计算同月度同城市均值
    mon_city_means = result.dropna(subset=['月度', '职级']).groupby(['月度', '城市名称', '职级']).apply(
        lambda x: x['发票含税金额'].sum() / x['住宿天数'].sum()
    ).reset_index()
    mon_city_means.columns = ['月度', '城市名称', '职级', '同月度同城市均值']

    # 计算同一酒店均值
    hotel_means = result.groupby(['年度', '城市名称', '职级', '宾馆']).apply(
        lambda x: x['发票含税金额'].sum() / x['住宿天数'].sum()
    ).reset_index()
    hotel_means.columns = ['年度', '城市名称', '职级', '宾馆', '同年度同酒店均值']

    # 计算同月度同一酒店均值
    mon_hotel_means = result.dropna(subset=['月度']).groupby(['月度', '城市名称', '职级', '宾馆']).apply(
        lambda x: x['发票含税金额'].sum() / x['住宿天数'].sum()
    ).reset_index()
    mon_hotel_means.columns = ['月度', '城市名称', '职级', '宾馆', '同月度同酒店均值']

    # 计算每天的住宿费
    result['每天住宿费'] = result['发票含税金额'] / result['住宿天数']
    result["每天发票含税金额"] = result['发票含税金额'] / result['住宿天数']

    # 合并均值数据到原始数据
    result = result.merge(mon_hotel_means, on=['月度', '城市名称', '宾馆', '职级'], how='left')
    result['同月度同酒店住宿费偏差率_数值'] = (result['每天住宿费'] - result['同月度同酒店均值']) / result['同月度同酒店均值']
    result['同月度同酒店住宿费偏差率'] = result['同月度同酒店住宿费偏差率_数值'].apply(lambda x: f"{x:.2%}" if pd.notnull(x) else None)

    result = result.merge(mon_city_means, on=['月度', '城市名称', '职级'], how='left')
    result['同月度同城市住宿费偏差率_数值'] = (result['每天住宿费'] - result['同月度同城市均值']) / result['同月度同城市均值']
    result['同月度同城市住宿费偏差率'] = result['同月度同城市住宿费偏差率_数值'].apply(lambda x: f"{x:.2%}" if pd.notnull(x) else None)

    result = result.merge(hotel_means, on=['年度', '城市名称', '宾馆', '职级'], how='left')
    result['同年度同酒店住宿费偏差率_数值'] = (result['每天住宿费'] - result['同年度同酒店均值']) / result['同年度同酒店均值']
    result['同年度同酒店住宿费偏差率'] = result['同年度同酒店住宿费偏差率_数值'].apply(lambda x: f"{x:.2%}" if pd.notnull(x) else None)

    result = result.merge(city_means, on=['年度', '城市名称', '职级'], how='left')
    result['同年度同城市住宿费偏差率_数值'] = (result['每天住宿费'] - result['同年度同城市均值']) / result['同年度同城市均值']
    result['同年度同城市住宿费偏差率'] = result['同年度同城市住宿费偏差率_数值'].apply(lambda x: f"{x:.2%}" if pd.notnull(x) else None)

    # 根据维度选择使用哪个偏差率和差价
    # 计算同城市差价（每天）
    result['同年度同城市差价'] = result['每天住宿费'] - result['同年度同城市均值']
    result['同月度同城市差价'] = result['每天住宿费'] - result['同月度同城市均值']
    # 使用同城市偏差率计算档位
    result['同年度同城市档位'] = result['同年度同城市住宿费偏差率_数值'].apply(lambda x: get_level(x) if pd.notnull(x) else None)
    result['同月度同城市档位'] = result['同月度同城市住宿费偏差率_数值'].apply(lambda x: get_level(x) if pd.notnull(x) else None)
    # 过滤掉包含inf的行
    result = result[~result['同年度同城市住宿费偏差率'].str.contains('inf', na=False)]
    result = result[~result['同月度同城市住宿费偏差率'].str.contains('inf', na=False)]

    # 计算同一酒店差价（每天）
    result['同年度同酒店差价'] = result['每天住宿费'] - result['同年度同酒店均值']
    result['同月度同酒店差价'] = result['每天住宿费'] - result['同月度同酒店均值']
    # 使用同一酒店偏差率计算档位
    result['同年度同酒店档位'] = result['同年度同酒店住宿费偏差率_数值'].apply(lambda x: get_level(x) if pd.notnull(x) else None)
    result['同月度同酒店档位'] = result['同月度同酒店住宿费偏差率_数值'].apply(lambda x: get_level(x) if pd.notnull(x) else None)
    # 过滤掉负偏差率的数据

    # 过滤掉包含inf的行
    result = result[~result['同年度同酒店住宿费偏差率'].str.contains('inf', na=False)]
    result = result[~result['同月度同酒店住宿费偏差率'].str.contains('inf', na=False)]


    if dimension == ['同城市']:
        result = result[(result['同年度同城市住宿费偏差率_数值'] > 0) | (result['同月度同城市住宿费偏差率_数值'] > 0)]
        result = result.drop(['同年度同酒店差价', '同月度同酒店差价', '同年度同酒店档位','同月度同酒店档位'], axis=1)
    elif dimension == ['同一酒店']:
        result = result[(result['同年度同酒店住宿费偏差率_数值'] > 0) | (result['同月度同酒店住宿费偏差率_数值'] > 0)]
        result = result.drop(['同年度同城市差价', '同月度同城市差价', '同年度同城市档位', '同月度同城市档位'], axis=1)
    elif len(dimension) == 2:
        result = result[(result['同年度同城市住宿费偏差率_数值'] > 0) | (result['同月度同城市住宿费偏差率_数值'] > 0) | (result['同年度同酒店住宿费偏差率_数值'] > 0) | (result['同月度同酒店住宿费偏差率_数值'] > 0)]

    # 删除没有发票号码
    result.dropna(subset=['发票号码'], inplace = True)
    result = result.round(2)

    # 删除中间计算列
    result = result.drop(['同年度同城市住宿费偏差率_数值', '同年度同酒店住宿费偏差率_数值', '每天住宿费',
                          '同月度同城市住宿费偏差率_数值', '同月度同酒店住宿费偏差率_数值'], axis=1)

    result = result.drop_duplicates(keep='first')
    result = result.sort_values(by = "同月度同酒店住宿费偏差率", ascending=False)

    return result

def get_level(rate):
    """
    根据偏差率计算档位
    
    Parameters:
    rate: float - 偏差率
    """
    if pd.isna(rate):
        return None
    # 对偏差率取绝对值后再判断档位
    abs_rate = abs(rate) * 100  # 转换为百分比
    if abs_rate <= 20:
        return '一档0-20%'
    elif abs_rate <= 40:
        return '二档20-40%'
    elif abs_rate <= 60:
        return '三档40-60%'
    elif abs_rate <= 80:
        return '四档60-80%'
    elif abs_rate <= 100:
        return '五档80-100%'
    else:
        return '六档≥100%'

def filter_by_level(df, selected_levels, dimension, month_year):
    """
    按档位筛选数据
    
    Parameters:
    df: DataFrame - 计算完指标的数据
    selected_levels: list - 选择的档位列表
    """
    if dimension == ["同城市"]:
        if selected_levels:
            if month_year == ["月度"]:
            # 同时筛选年度档位和月度档位
                return df[df['同月度同城市档位'].isin(selected_levels)]
            elif month_year == ["年度"]:
            # 同时筛选年度档位和月度档位
                return df[df['同年度同城市档位'].isin(selected_levels)]
            elif len(month_year) == 2:
                return df[(df['同月度同城市档位'].isin(selected_levels)) | (df['同年度同城市档位'].isin(selected_levels))]
    elif dimension == ["同一酒店"]:
        if selected_levels:
            if month_year == ["月度"]:
            # 同时筛选年度档位和月度档位
                return df[df['同月度同酒店档位'].isin(selected_levels)]
            elif month_year == ["年度"]:
            # 同时筛选年度档位和月度档位
                return df[df['同年度同酒店档位'].isin(selected_levels)]
            elif len(month_year) == 2:
                return df[(df['同月度同酒店档位'].isin(selected_levels)) | (df['同年度同酒店档位'].isin(selected_levels))]
    elif len(dimension) == 2:
        if selected_levels:
            if month_year == ["月度"]:
            # 同时筛选年度档位和月度档位
                return df[(df['同月度同酒店档位'].isin(selected_levels)) | (df['同月度同城市档位'].isin(selected_levels))]
            elif month_year == ["年度"]:
            # 同时筛选年度档位和月度档位
                return df[(df['同年度同酒店档位'].isin(selected_levels)) | (df['同年度同城市档位'].isin(selected_levels))]
            elif len(month_year) == 2:
                return df[(df['同月度同酒店档位'].isin(selected_levels)) | (df['同年度同酒店档位'].isin(selected_levels)) | (df['同月度同城市档位'].isin(selected_levels)) | (df['同年度同城市档位'].isin(selected_levels))]

    return df

# 示例使用
if __name__ == "__main__":
    # 读取Excel文件
    df = pd.read_excel("data/差旅费票据信息.xlsx")
    
    # 示例：计算2023年的数据
    selected_years = [2023]
    dimension = ['同一酒店']  # 或 '同一酒店'
    selected_levels = ['一档0-20%', '二档20-40%']
    month_year = ['月度', '年度']  # 新增：级别选择
    
    # 计算指标
    result = calculate_metrics(df, selected_years, dimension)
    
    # 按档位筛选
    filtered_result = filter_by_level(result, selected_levels, dimension, month_year)
    
    # 输出结果到Excel
    filtered_result.to_excel("分析结果.xlsx", index=False)
    print("分析完成，结果已保存到分析结果.xlsx")